#include "VideoEncoder.h"

VideoEncoder::VideoEncoder()
{

}

VideoEncoder::~VideoEncoder()
{
    this->DeInit();
}

int VideoEncoder::InitH264(int width, int height, int fps, int bit_rate)
{
    width_ = width;
    height_ = height;
    fps_ = fps;
    bit_rate_ = bit_rate;

    AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (codec == NULL) {
        printf("avcodec_find_encoder AV_CODEC_ID_H264 failed.\n");
        return -1;
    }

    codec_ctx_ = avcodec_alloc_context3(codec);
    if (codec_ctx_ == NULL) {
        printf("avcodec_alloc_context3 AV_CODEC_ID_H264 failed.\n");
        return -1;
    }

    codec_ctx_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
    codec_ctx_->bit_rate = bit_rate_;
    codec_ctx_->width = width_;
    codec_ctx_->height = height_;
    codec_ctx_->framerate = {fps_, 1};
    //不设置的话会报错"The encoder time_base is not set."
    codec_ctx_->time_base = {1, 1000 * 1000}; //单位为us
    codec_ctx_->gop_size = fps_;
    codec_ctx_->max_b_frames = 0;
    codec_ctx_->pix_fmt = AV_PIX_FMT_YUV420P;
    //编码器延迟太严重，通过这个设置可以降低延迟
    av_dict_set(&dict_, "tune", "zerolatency", 0);

    int ret = avcodec_open2(codec_ctx_, NULL, &dict_);
    if (ret != 0) {
        printf("avcodec_open2 failed:%s\n", printErrStr(ret));
        return -1;
    }

    frame_ = av_frame_alloc();
    if (frame_ == NULL) {
        printf("av_frame_alloc failed.\n");
        return -1;
    }

    frame_->width = width_;
    frame_->height = height_;
    frame_->format = codec_ctx_->pix_fmt;

    printf("Init H264 success.\n");
    return 0;
}

void VideoEncoder::DeInit()
{
    if (codec_ctx_ != NULL) {
        avcodec_free_context(&codec_ctx_);
    }

    if (frame_ != NULL)
        av_frame_free(&frame_);

    if (dict_ != NULL)
        av_dict_free(&dict_);
}

AVPacket *VideoEncoder::Encode(uint8_t *yuv_data, int yuv_size, int stream_index, int64_t pts, int64_t time_base)
{
    if (codec_ctx_ == NULL) {
        printf("codec_ctx is NULL.\n");
        return NULL;
    }

    int ret = 0;
    pts = av_rescale_q(pts, AVRational{1, (int)time_base}, codec_ctx_->time_base);
    frame_->pts = pts;
    if (yuv_data != NULL) {
        int ret_size = av_image_fill_arrays(frame_->data, frame_->linesize, yuv_data, (AVPixelFormat)frame_->format, frame_->width, frame_->height, 1);

        if (ret_size != yuv_size) {
            printf("ret_size:%d != yuv_size:%d -> failed.\n", ret_size, yuv_size);
            return NULL;
        }
        ret = avcodec_send_frame(codec_ctx_, frame_);
    }
    else {
        ret = avcodec_send_frame(codec_ctx_, NULL);
    }

    if (ret != 0) {
        printf("avcodec_send_frame failed:%s\n", printErrStr(ret));
        return NULL;
    }

    AVPacket *packet = av_packet_alloc();
    ret = avcodec_receive_packet(codec_ctx_, packet);
    if (ret != 0) {
        printf("h264 avcodec_receive_packet failed:%s\n", printErrStr(ret));
        return NULL;
    }

    packet->stream_index = stream_index;
    return packet;
}

int VideoEncoder::Encode(uint8_t *yuv_data, int yuv_size, int stream_index, int64_t pts, int64_t time_base, std::vector<AVPacket *> &packets)
{
    if (codec_ctx_ == NULL) {
        printf("codec_ctx is NULL.\n");
        return -1;
    }

    int ret = 0;
    pts = av_rescale_q(pts, AVRational{1, (int)time_base}, codec_ctx_->time_base);
    frame_->pts = pts;
    if (yuv_data != NULL) {
        int ret_size = av_image_fill_arrays(frame_->data, frame_->linesize, yuv_data, (AVPixelFormat)frame_->format, frame_->width, frame_->height, 1);

        if (ret_size != yuv_size) {
            printf("ret_size:%d != yuv_size:%d -> failed.\n", ret_size, yuv_size);
            return -1;
        }
        ret = avcodec_send_frame(codec_ctx_, frame_);
    }
    else {
        ret = avcodec_send_frame(codec_ctx_, NULL);
    }

    if (ret != 0) {
        printf("avcodec_send_frame failed:%s\n", printErrStr(ret));
        return -1;
    }

    while (1) {
        AVPacket *packet = av_packet_alloc();
        ret = avcodec_receive_packet(codec_ctx_, packet);
        packet->stream_index = stream_index;
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            ret = 0;
            av_packet_free(&packet);
            break;
        }
        else if(ret < 0) {
            printf("avcodec_receive_packet h264 failed:%s\n", printErrStr(ret));
            av_packet_free(&packet);
            ret = -1;
        }
        printf("h264 pts:%lld\n", packet->pts);
        packets.push_back(packet);
    }
    return ret;
}

AVCodecContext *VideoEncoder::GetCodecContext()
{
    return codec_ctx_;
}

char* VideoEncoder::printErrStr(int errNo)
{
    char tmp[AV_ERROR_MAX_STRING_SIZE] = {0};
    return av_make_error_string(tmp, AV_ERROR_MAX_STRING_SIZE, errNo);
}




